
(* TIM *)

include Functional

type instruction =
	| Take of int
	| Enter of addrMode
	| Push of addrMode
and addrMode =
	| Arg of int
	| Label of string
	| Code of instruction list
	| IntConst of int

type framePointer =
	| FrameAddr of int
	| FrameInt of int
	| FrameNull

type closure = instruction list * framePointer

type stack = closure list

type frame = closure list

type codeStore = string * instruction list

let fail fmt = Printf.kprintf failwith fmt

let codeLookup cstore label =
	try
		List.assoc label cstore
	with exn ->
		fail "Jump to unknown label %S" label

let preludeDefs = Functional.prelude
let compiledPrimitives = []

let rec zipArgs_aux n = function
	| [] -> []
	| a :: t -> (a, Arg n) :: zipArgs_aux (n+1) t

let zipArgs args = zipArgs_aux 1 args

let rec compileR expr env = match expr with
	| EAp (e1, e2) ->
		Push (compileA e2 env) :: compileR e1 env
	| EVar _ (* same as ENum *)
	| ENum _ ->
		[Enter (compileA expr env)]
	| _ -> fail "compileR: not fancy enough yet :P"

and compileA expr env = match expr with
	| EVar v -> begin
			try
				List.assoc v env
			with Not_found ->
				fail "Unknown variable %S" v
		end
	| ENum n -> IntConst n
	| _ -> Code (compileR expr env)

let compileSC env (name, args, body) =
	let newEnv = zipArgs args @ env in
	let insns = compileR body newEnv in
	(name, Take (List.length args) :: insns)

let compile program =
	let defns = preludeDefs @ program in
	let env =
		(List.map (fun (name, _, body) -> (name, Label name)) defns)
		@
		(List.map (fun (name, _) -> (name, Label name)) compiledPrimitives)
	in
	let compiledDefns = List.map (compileSC env) defns in
	let compiledCode = compiledDefns @ compiledPrimitives in
	(
		[Enter (Label "mymain")],
		FrameNull,
		[],
		compiledCode
	)

let main = ("mymain", [], ap (ap (ap (var "S") (var "K")) (var "K")) (num 4))

(* The evaluator/interpreter *)

(* blargh! :P *)

(* So compile gives us our bytecode... *)
(* Now we need to turn it into machine code :P *)

(* Heh, the Mark 1 evaluator is stupidly simple! *)
(* A task for tomorrow, eh? ;-) *)

(*

Take <n>
	modifies: stack, frame pointer, heap
	tuple, heap = alloc (N from stack) heap
	frame pointer = tuple
	stack = remove N from stack

Enter <am>:
	modifies: program counter, frame pointer
	program counter, frame pointer = amToClosure AM

Push <am>:
	modifies: stack
	stack = amToClosure AM

stack contains: closures
closure: code + frame pointer value
closure: program counter + frame pointer

frame pointer could just point into the stack I guess... hmm. but
then we never free up the stack...

amToClosure:

Arg <n>:
	N is index into the frame pointer

Code <il>:
	IL is the code address to jump to
	(this is will be code emitted inside the function's definition to
	preserve locality I guess)

Label <l>:
	L is the code address to jump to (heh)

IntConst <n>:
	N is a simple constant (char/bool/int/etc)

*)

(*

Machine State:

stack pointer (esp), frame pointer (ebp), accumulator (eax)

closure = (code, frame)

code = machine address (in code segment)
frame = NULL | machine address | integer constant

*)

let pp fmt = Printf.printf (fmt ^^ format_of_string "\n")

let ppt fmt = pp (format_of_string "\t" ^^ fmt)

let num =
	let next = ref 0 in
	(fun () -> incr next; !next)

let rec emit (init,_,_,defns) =
	emitInit init;
	List.iter emitDefn defns

and emitInit = function
	(* just one slightly small problem :) I screw up the C stack layout *)
	| [Enter (Label "mymain")] ->
		pp ".section .data";
		pp ".global _timstart";
		pp ".section .text";
		pp "_intCode:";
		ppt "movl %%ebp, %%eax";
		ppt "jmp _timfinish";
		pp "_timstart:";
		ppt "movl %%esp, %%ebx";
		ppt "call mymain";
		pp "_timfinish:";
		ppt "movl %%ebx, %%esp";
		ppt "ret"
	| _ -> failwith "Invalid prologue"

and emitDefn (name, insns) =
	pp "%s:" name;
	emitInsns insns

(* fairly fucking lame :P but it works... *)
and emitInsns = function
	| Take 0 :: rest ->
		emitInsns rest
	| Take n :: rest ->
		(* take 2n machine words from the stack, allocate into a frame *)
		(* put frame into frame pointer *)
		(* add 2n to stack pointer *)
		
		(* -- currently use stack as heap, and never free stack *)
		ppt "movl %%esp, %%ebp";
		emitInsns rest
	| Push (IntConst n) :: rest ->
		ppt "pushl $%d" n;
		ppt "push $%s" "_intCode";
		emitInsns rest
	| Push (Label lbl) :: rest ->
		ppt "push %%ebp";
		ppt "push $%s" lbl;
		emitInsns rest
	| Push (Arg n) :: rest ->
		ppt "push %d(%%ebp)" (n * 8 - 4);
		ppt "push %d(%%ebp)" (n * 8 - 8);
		emitInsns rest
	| Push (Code c) :: rest ->
		let label = Printf.sprintf "_PL%d" (num ()) in
		ppt "push %%ebp";
		ppt "push $%s" label;
		emitInsns rest;
		emitDefn (label, c)
	| Enter (IntConst n) :: rest ->
		ppt "movl $%d, %%ebp" n;
		ppt "jmp _intCode";
		emitInsns rest
	| Enter (Label lbl) :: rest ->
		(* no change to frame pointer *)
		ppt "jmp %s" lbl;
		emitInsns rest
	| Enter (Code c) :: rest ->
		(* no change to frame pointer *)
		let label = Printf.sprintf "_EL%d" (num ()) in
		ppt "jmp %s" label;
		emitInsns rest;
		emitDefn (label, c)
	| Enter (Arg n) :: rest ->
		(* want the nth tuple from the current frame pointer *)
		(* can't update frame pointer until we've gotten everything out of it *)
		ppt "movl %d(%%ebp), %%ecx" (n * 8 - 4);
		ppt "movl %d(%%ebp), %%edx" (n * 8 - 8);
		ppt "movl %%ecx, %%ebp";
		ppt "jmp *%%edx";
		emitInsns rest
	| [] ->
		(* done *)
		()

(* amToClosure:
	Arg n -> fGet heap fptr n
	Code il -> il, fptr
	Label l -> codeLookup cstore l, fptr
	IntConst n -> intCode, FrameInt n
*)
;;

(*
	cons a b cn cc = cc a b
	nil cn cc = cn
	caseList = I
	inf x = cons x (inf x)
	ones = inf 1
	hd xs = caseList xs 0 hd_cons
	hd_cons x xs = x
*)

let num = Functional.num

let program = [
	("cons", ["a"; "b"; "cn"; "cc"], (ap (ap (var "cc") (var "a")) (var "b")));
	("nil", ["cn"; "cc"], (var "cn"));
	("caseList", [], (var "I"));
	("inf", ["x"], (ap (ap (var "cons") (var "x")) (ap (var "inf") (var "x"))));
	("ones", [], (ap (var "inf") (num 1)));
	("hd", ["xs"], (ap (ap (ap (var "caseList") (var "xs")) (num 0)) (var "hd_cons")));
	("hd_cons", ["x"; "xs"], (var "x"));
]

let main = ("mymain", [], (ap (var "hd") (var "ones")))

let _ = emit (compile (main::program))
